package org.mockserver.proxy.relay;
import com.google.common.annotations.VisibleForTesting;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.HttpClientCodec;
import io.netty.handler.codec.http.HttpContentDecompressor;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.ssl.SslHandler;
import org.mockserver.logging.LoggingHandler;
import org.mockserver.proxy.http.HttpProxy;
import org.mockserver.proxy.unification.PortUnificationHandler;
import org.mockserver.socket.SSLFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.net.InetSocketAddress;
import static org.mockserver.proxy.error.Logging.shouldIgnoreException;
@ChannelHandler.Sharable
public abstract class RelayConnectHandler<T> extends SimpleChannelInboundHandler<T> {
@VisibleForTesting
public static Logger logger = LoggerFactory.getLogger(RelayConnectHandler.class);
@Override
public void channelRead0(final ChannelHandlerContext serverCtx, final T request) throws Exception {
Bootstrap bootstrap = new Bootstrap()
.group(serverCtx.channel().eventLoop())
.channel(NioSocketChannel.class)
.handler(new ChannelInboundHandlerAdapter() {
@Override
public void channelActive(final ChannelHandlerContext clientCtx) throws Exception {
serverCtx.channel()
.writeAndFlush(successResponse(request))
.addListener(new ChannelFutureListener() {
@Override
public void operationComplete(ChannelFuture channelFuture) throws Exception {
removeCodecSupport(serverCtx);
// downstream
ChannelPipeline downstreamPipeline = clientCtx.channel().pipeline();
if (PortUnificationHandler.isSslEnabledDownstream(serverCtx.channel())) {
downstreamPipeline.addLast(new SslHandler(SSLFactory.createClientSSLEngine()));
}
if (logger.isDebugEnabled()) {
downstreamPipeline.addLast(new LoggingHandler(this.getClass().getSimpleName() + " -->"));
}
downstreamPipeline.addLast(new HttpClientCodec());
downstreamPipeline.addLast(new HttpContentDecompressor());
downstreamPipeline.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
downstreamPipeline.addLast(new DownstreamProxyRelayHandler(serverCtx.channel(), logger));
// upstream
ChannelPipeline upstreamPipeline = serverCtx.channel().pipeline();
if (PortUnificationHandler.isSslEnabledUpstream(serverCtx.channel())) {
upstreamPipeline.addLast(new SslHandler(SSLFactory.createServerSSLEngine()));
}
if (logger.isDebugEnabled()) {
upstreamPipeline.addLast(new LoggingHandler(this.getClass().getSimpleName() + "<-- "));
}
upstreamPipeline.addLast(new HttpServerCodec());
upstreamPipeline.addLast(new HttpContentDecompressor());
upstreamPipeline.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
upstreamPipeline.addLast(new UpstreamProxyRelayHandler(clientCtx.channel(), logger));
}
});
}
});
final InetSocketAddress remoteSocket = getDownstreamSocket(serverCtx.channel());
bootstrap.connect(remoteSocket).addListener(new ChannelFutureListener() {
@Override
public void operationComplete(ChannelFuture future) throws Exception {
if (!future.isSuccess()) {
failure("Connection failed to " + remoteSocket, future.cause(), serverCtx, failureResponse(request));
}
}
});
}
private InetSocketAddress getDownstreamSocket(Channel channel) {
if (channel.attr(HttpProxy.REMOTE_SOCKET).get() != null) {
return channel.attr(HttpProxy.REMOTE_SOCKET).get();
} else if (channel.attr(HttpProxy.HTTP_CONNECT_SOCKET).get() != null) {
return channel.attr(HttpProxy.HTTP_CONNECT_SOCKET).get();
} else {
throw new IllegalStateException("Trying to connect to remote socket but no remote socket has been set");
}
}
@Override
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
failure("Exception caught by CONNECT proxy handler -> closing pipeline ", cause, ctx, failureResponse(null));
}
private void failure(String message, Throwable cause, ChannelHandlerContext ctx, Object response) {
if (!shouldIgnoreException(cause)) {
logger.warn(message, cause);
}
Channel channel = ctx.channel();
channel.writeAndFlush(response);
if (channel.isActive()) {
channel.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
}
}
protected abstract void removeCodecSupport(ChannelHandlerContext ctx);
protected abstract Object successResponse(Object request);
protected abstract Object failureResponse(Object request);
protected void removeHandler(ChannelPipeline pipeline, Class<? extends ChannelHandler> handlerType) {
if (pipeline.get(handlerType) != null) {
pipeline.remove(handlerType);
}
}
protected void removeHandler(ChannelPipeline pipeline, ChannelHandler channelHandler) {
if (pipeline.toMap().containsValue(channelHandler)) {
pipeline.remove(channelHandler);
}
}
}